from typing import List

import numpy as np

from .map_summarizer import MAPSummarizer


class ODSummarizer(MAPSummarizer):
    def __init__(
        self,
        iou_thresh: float,
        class_names: List[str],
        scale_splitter: List[float] = None,
    ):
        """
        :param iou_thresh:float
            交并比阈值，决定着prediction是否能够匹配ground truth
        :param class_names:List[str]
            类别名称，如车辆检测的["car", "bus", "truck"]
        :param scale_splitter:List[float]=None
            尺度分隔基准，如：将目标按照宽度分成三组([0, 20), [20, 200), [200, 512])，则此处应传入：[0, 20, 200, 512]
        """
        super().__init__(class_names)
        self.iou_thresh = iou_thresh
        self.scale_splitter = scale_splitter
        if scale_splitter:
            self.scales_names = [
                "(" + str(w_min) + "~" + str(w_max) + ")"
                for w_min, w_max in zip(scale_splitter[:-1], scale_splitter[1:])
            ]

    def get_frame_cache(self, result: np.ndarray, label: np.ndarray):
        """依据检测结果和真实值，获取并缓存当前帧的评估结果缓存

        :param result:np.ndarray
            结果数组，每行一个目标，列的排布顺序为：conf,x1,y1,x2,y2,class
        :param label:
            真值数组，每行一个目标，列的排布顺序为：x1,y1,x2,y2,class
        :return:
            缓存列表，每个元素为一个prediction的评估结果，结果格式为:Tuple[TPFlag, conf, class]
        """
        result = result[np.argsort(-result[:, 0])]  # 按照概率倒排
        iou_mat = self.cal_iou_mat(result[:, 1:5], label[:, :4])
        msk_cls = (np.expand_dims(result[:, -1], 1) - label[:, -1]) == 0  # 类别相同
        msk_iou = iou_mat >= self.iou_thresh  # iou足够大
        msk_match = iou_mat == iou_mat.max(1, keepdims=True)  # 最优匹配
        pred_id, gt_id = np.where(iou_mat * msk_cls * msk_iou * msk_match > 0)
        gt_id = gt_id.tolist()
        msk_first_match = [gt_id.index(v) == i for i, v in enumerate(gt_id)]
        tp_ids = pred_id[msk_first_match].tolist()
        conf = result[:, 0]
        cls = self.get_cls_cache(result[:, 1:])
        return [(i in tp_ids, *_) for i, _ in enumerate(zip(conf, cls))]

    def get_cls_cache(self, label):
        if len(label.shape) == 2:  # 标签非空，目标存在
            type_cls = label[:, -1]
            if self.scale_splitter:
                size_cls = self.get_size_cls(label)
                return np.stack([type_cls, size_cls], 1)
            return type_cls
        return np.empty([0, 2])

    def get_size_cls(self, label):
        widths = label[:, 2] - label[:, 0]
        if not self.scale_splitter:
            return np.zeros_like(widths)
        masks = []
        for size_cls, (w_min, w_max) in enumerate(
            zip(self.scale_splitter[:-1], self.scale_splitter[1:])
        ):
            mask = np.logical_and(widths > w_min, widths < w_max)
            masks.append(mask)
        return np.stack(masks).argmax(0)

    def calculate_metrics(self):
        metric_cache = np.concatenate(self.metric_cache)
        metric_cache = metric_cache[np.argsort(-metric_cache[:, 1])]
        gt_cls_cache = np.concatenate(self.gt_cls_cache)
        metrics = dict()
        ap_array = np.empty(len(self.class_names))
        ap_matrix = np.empty([len(self.class_names), len(self.scales_names)])
        gt_type, gt_scale = gt_cls_cache.transpose()
        p_type, p_scale = np.array(metric_cache[:, 2].tolist()).transpose()
        for cls_idx, cls_name in enumerate(self.class_names):
            gt_type_msk = gt_type == cls_idx
            p_type_msk = p_type == cls_idx
            ap = self.calculate_ap(
                total_objects=gt_type_msk.sum(),
                metric_cache=metric_cache[p_type_msk, 0],
            )
            ap_array[cls_idx] = ap
            metrics[cls_name + "AP"] = ap
            for scale_idx, scale_name in enumerate(self.scales_names):
                ap = self.calculate_ap(
                    total_objects=(gt_type_msk * (gt_scale == scale_idx)).sum(),
                    metric_cache=metric_cache[p_type_msk * (p_scale == scale_idx), 0],
                )
                ap_matrix[cls_idx, scale_idx] = ap
                metrics[cls_name + "AP" + scale_name] = ap

        metrics["mAP"] = ap_array.mean()
        for scale_map, scale_name in zip(ap_matrix.mean(0), self.scales_names):
            metrics["AP" + scale_name] = scale_map
        return metrics

    @staticmethod
    def cal_iou_mat(box1, box2):
        # Get the coordinates of bounding boxes
        x11, y11, x12, y12 = box1.transpose()
        x21, y21, x22, y22 = np.expand_dims(box2, -1).transpose([1, 0, 2])
        # get the coordinates of the intersection rectangle
        wi = np.maximum(np.minimum(x12, x22) - np.maximum(x11, x21), 0)
        hi = np.maximum(np.minimum(y12, y22) - np.maximum(y11, y21), 0)
        si = wi * hi
        su = (x12 - x11) * (y12 - y11) + (x22 - x21) * (y22 - y21) - si
        iou_mat = si / su
        return iou_mat.transpose()
